package at.chrenko.tu.vs.shared;

import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

public class RSADecorator extends ChannelDecorator {

	private PrivateKey privateKey;
	private PublicKey publicKey;
	private Cipher crypt;
	
	public RSADecorator(Channel channel, PrivateKey privateKey, PublicKey publicKey) throws NoSuchAlgorithmException, NoSuchPaddingException {
		super(channel);
		this.privateKey = privateKey;
		this.publicKey = publicKey;
		
		this.crypt = Cipher.getInstance("RSA/NONE/OAEPWithSHA256AndMGF1Padding");
	}
	
	@Override
	public void send(byte[] message) throws IOException {
		if(message == null)
			return;
		
		try {
			crypt.init(Cipher.ENCRYPT_MODE, publicKey);
		} catch (InvalidKeyException e) {
			System.err.println(e.getMessage());
			return;
		}
		
		try {
			decoratedChannel.send(crypt.doFinal(message));
		} catch (IllegalBlockSizeException e) {
			System.err.println(e.getMessage());
			return;
		} catch (BadPaddingException e) {
			System.err.println(e.getMessage());
			return;
		}
	}
	
	@Override
	public byte[] receive() throws IOException {
		byte[] data = decoratedChannel.receive();
		
		if(data == null)
			return null;
		
		try {
			crypt.init(Cipher.DECRYPT_MODE, privateKey);
		} catch (InvalidKeyException e) {
			System.err.println(e.getMessage());
			return null;
		}
		
		try {
			return crypt.doFinal(data);
		} catch (IllegalBlockSizeException e) {
			System.err.println(e.getMessage());
			return null;
		} catch (BadPaddingException e) {
			System.err.println(e.getMessage());
			return null;
		}
	}
}
